package com.ruoyi.common.currentlimit.aspect;

import com.ruoyi.common.currentlimit.annotation.RedisLimitAnnotation;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;
import java.util.Collections;
import java.util.List;

/**
 * @description: 限流切面类
 * @author: Yuan
 * @create: 2025-02-22 08:54
 **/
@Aspect
@Component
@Slf4j
public class LimitAspect {

    @Autowired
    private RedisTemplate<String ,Object> redisTemplate;

    @Autowired
    private DefaultRedisScript<Long> redisScript;

    @Pointcut(value = "@annotation(com.ruoyi.common.currentlimit.annotation.RedisLimitAnnotation)")
    public void rateLimit(){

    }

    @Around("rateLimit()")
    public Object interceptor(ProceedingJoinPoint joinPoint) throws Throwable {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        Class<?> targetClass = method.getDeclaringClass();
        RedisLimitAnnotation rateLimit = method.getAnnotation(RedisLimitAnnotation.class);
        if (rateLimit != null){
            HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();
            String ipAddr = getIpAddr(request);
            StringBuffer stringBuffer = new StringBuffer();
            stringBuffer.append(ipAddr).append("-")
                    .append(targetClass.getName()).append("-")
                    .append(method.getName()).append("-")
                    .append(rateLimit.key());
            List<String> keys = Collections.singletonList(stringBuffer.toString());
            Number number = redisTemplate.execute(redisScript, keys, rateLimit.count(), rateLimit.period());
            if (number != null && number.intValue() != 0 && number.intValue() <= rateLimit.count()){
                log.info("限流时间段内访问了第:{}次",number.toString());
                return joinPoint.proceed();
            }
        }else {
            return joinPoint.proceed();
        }
        throw new RuntimeException("访问频率过快，被限流了");
    }

    /***
    * @Description: 获取请求ip
    * @Param: [request]
    * @return: java.lang.String
    * @Author: Yuan
    * @Date: 2025/2/22
    */
    public static String getIpAddr(HttpServletRequest request){
        String ipAddress = null;
        try {
            ipAddress = request.getHeader("x-forwarder-for");
            if (ipAddress == null || ipAddress.length() == 0 || "unknown".equalsIgnoreCase(ipAddress)){
                ipAddress = request.getHeader("Proxy-Client-IP");
            }
            if (ipAddress == null || ipAddress.length() == 0 || "unknown".equalsIgnoreCase(ipAddress)){
                ipAddress = request.getHeader("WL-Proxy-Client-IP");
            }
            if (ipAddress == null || ipAddress.length() == 0 || "unknown".equalsIgnoreCase(ipAddress)){
                ipAddress = request.getRemoteAddr();
            }
            //对于通过多个代理的情况，第一个ip为真实ip，多个ip按照','分割
            if (ipAddress != null && ipAddress.length() > 15){
                if (ipAddress.indexOf(",") > 0){
                    ipAddress = ipAddress.substring(0,ipAddress.indexOf(","));
                }
            }
        } catch (Exception e) {
            ipAddress = "";
        }
        return ipAddress;
    }

}
